import numpy as np
from PIL import Image
from torch.utils.tensorboard import SummaryWriter


class Logger:
    def __init__(self, log_dir):
        """Initialize SummaryWriter."""
        self.writer = SummaryWriter(log_dir)

    def scalar_summary(self, tag, value, step):
        """Log a scalar variable."""
        self.writer.add_scalar(tag, value, step)

    def image_summary(self, tag, image, step):
        """Log a single image."""
        image = Image.fromarray((image * 255).astype(np.uint8))
        self.writer.add_image(tag, np.array(image), step, dataformats="HWC")

    def image_list_summary(self, tag, images, step):
        """Log a list of images."""
        for i, img in enumerate(images):
            image = Image.fromarray((img * 255).astype(np.uint8))
            self.writer.add_image(
                f"{tag}/{i}", np.array(image), step, dataformats="HWC"
            )

    def close(self):
        """Close the SummaryWriter."""
        self.writer.close()
